package cn.pengpeng;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.mapreduce.*;
import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter;
import org.apache.log4j.Logger;

import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

public class MySQLOutputFormat<V extends MySQLOutputWritable> extends OutputFormat<NullWritable, V> {

    private static final Logger LOG = Logger.getLogger(MySQLOutputFormat.class);

    /** 配置 - 输出端数据库驱动类 */
    public static final String MYSQL_OUTPUT_DRIVER = "mysql.output.dirver";
    /** 配置 - 输出端数据库URL */
    public static final String MYSQL_OUTPUT_URL = "mysql.output.url";
    /** 配置 - 输出端数据库用户名 */
    public static final String MYSQL_OUTPUT_USERNAME = "mysql.output.username";
    /** 配置 - 输出端数据库密码 */
    public static final String MYSQL_OUTPUT_PASSWORD = "mysql.output.password";
    /** 配置 - 批量提交的数据记录数 */
    public static final String MYSQL_OUTPUT_BATCH_SIZE = "mysql.output.batch.size";

    /**
     * 获取记录写入对象
     */
    @Override
    public RecordWriter<NullWritable, V> getRecordWriter(TaskAttemptContext context)
            throws IOException, InterruptedException {

        return new MySQLRecordWriter(context.getConfiguration());
    }

    /**
     * 检查输出空间是否有效
     */
    @Override
    public void checkOutputSpecs(JobContext context)
            throws IOException, InterruptedException {
        Connection conn = null;
        try {
            conn = this.getConnection(context.getConfiguration());
        } catch (Exception e) {
            throw new IOException("连接数据库失败", e);
        } finally {
            this.closeAutoCloseable(conn);
        }
    }

    @Override
    public OutputCommitter getOutputCommitter(TaskAttemptContext context)
            throws IOException, InterruptedException {

        return new FileOutputCommitter(null, context);
    }


    /**
     * MySQL数据写入类
     */
    public class MySQLRecordWriter extends RecordWriter<NullWritable, V> {
        protected Configuration conf;
        private Connection conn;
        // PreparedStatement 缓冲器
        private Map<String, PreparedStatement> pstmtCache = new HashMap<String, PreparedStatement>();
        // Batch计数器
        private Map<String, Integer> batchCache = new HashMap<String, Integer>();
        private int batchSize = 100; // 批量提交记录数

        public MySQLRecordWriter() {}

        public MySQLRecordWriter(Configuration conf) {
            this.conf = conf;
            this.batchSize = conf.getInt(MYSQL_OUTPUT_BATCH_SIZE, this.batchSize);
        }

        @Override
        public void write(NullWritable key, V value) throws IOException, InterruptedException {

            if(this.conn == null) {
                try {
                    this.conn = MySQLOutputFormat.this.getConnection(this.conf);
                    this.conn.setAutoCommit(false); // 关闭自动提交
                } catch (Exception e) {
                    throw new IOException("连接数据库失败", e);
                }
            }

            String sql = value.fetchInsertOrUpdateSql();
            PreparedStatement pstmt = this.pstmtCache.get(sql);
            if(pstmt == null) {
                try {
                    pstmt = conn.prepareStatement(value.fetchInsertOrUpdateSql());
                    this.pstmtCache.put(sql, pstmt);
                } catch (SQLException e) {
                    throw new IOException("创建PreparedStatement对象产生异常", e);
                }
            }

            Integer count = this.batchCache.get(sql);
            if(count == null)
                count = 0;

            try {
                value.setPreparedStatementParameters(pstmt);
                pstmt.addBatch();
                count++;
                if(count >= this.batchSize) {
                    pstmt.executeBatch(); // 批量执行
                    this.conn.commit(); // 提交执行结果
                    count = 0; // 清零
                }
                this.batchCache.put(sql, count);
            } catch (SQLException e) {
                throw new IOException("向数据库写入数据出现异常", e);
            }
        }

        @Override
        public void close(TaskAttemptContext context) throws IOException, InterruptedException {

            // 将缓冲器中的pstmt再次提交一次，防止因批量提交数量不足而未提交的数据漏掉
            for(Map.Entry<String, PreparedStatement> entry : pstmtCache.entrySet()) {
                try {
                    entry.getValue().executeBatch();
                    this.conn.commit();
                } catch (SQLException e) {
                    throw new IOException("向数据库写入数据出现异常", e);
                }
            }

            MySQLOutputFormat.this.closeAutoCloseable(this.conn);
        }
    }

    /**
     * 获取数据库连接
     * @param conf
     * @return
     * @throws Exception
     */
    private Connection getConnection(Configuration conf) throws Exception {

        String driver = conf.get(MYSQL_OUTPUT_DRIVER);
        String url = conf.get(MYSQL_OUTPUT_URL);
        String username = conf.get(MYSQL_OUTPUT_USERNAME);
        String password = conf.get(MYSQL_OUTPUT_PASSWORD);

        Class.forName(driver);
        return DriverManager.getConnection(url, username, password);
    }

    /**
     * 关闭连接
     * @param autoCloseable
     */
    private void closeAutoCloseable(AutoCloseable autoCloseable) {
        try {
            if(autoCloseable != null)
                autoCloseable.close();
        } catch (Exception e) {
            LOG.error("关闭失败"+e.getMessage());
        }
    }
}
